package com.proxy.socks5;

import io.netty.bootstrap.Bootstrap;
import io.netty.channel.*;
import io.netty.channel.socket.SocketChannel;
import io.netty.channel.socket.nio.NioSocketChannel;
import io.netty.handler.codec.http.HttpServerCodec;
import io.netty.handler.codec.socksx.v5.*;
import io.netty.util.ReferenceCountUtil;

import java.util.Set;

/**
 * @author kdyzm
 * @date 2021-04-23
 */
public class Socks5CommandRequestInboundHandler extends SimpleChannelInboundHandler<DefaultSocks5CommandRequest> {

    private EventLoopGroup eventExecutors;

    private Set<String> blackList;

    public Socks5CommandRequestInboundHandler(EventLoopGroup eventExecutors, Set<String> blackList){
        this.eventExecutors = eventExecutors;
        this.blackList = blackList;
    }

    @Override
    protected void channelRead0(ChannelHandlerContext ctx, DefaultSocks5CommandRequest msg) throws Exception {
        Socks5AddressType socks5AddressType = msg.dstAddrType();
        if (!msg.type().equals(Socks5CommandType.CONNECT)) {
            System.out.println("receive commandRequest type=" +  msg.type());
            ReferenceCountUtil.retain(msg);
            ctx.fireChannelRead(msg);
            return;
        }
        //检查黑名单
        if (inBlackList(msg.dstAddr())) {
            System.out.println(msg.dstAddr() + " 地址在黑名单中，拒绝连接");
            //假装连接成功
            DefaultSocks5CommandResponse commandResponse = new DefaultSocks5CommandResponse(Socks5CommandStatus.SUCCESS, socks5AddressType);
            ctx.writeAndFlush(commandResponse);
            ctx.pipeline().addLast("HttpServerCodec",new HttpServerCodec());
            ctx.pipeline().addLast(new BlackListInboundHandler());
            ctx.pipeline().remove(Socks5CommandRequestInboundHandler.class);
            ctx.pipeline().remove(Socks5CommandRequestDecoder.class);
            return;
        }
        System.out.println("准备连接目标服务器，ip="+msg.dstAddr()+",port=" +  msg.dstPort());
        Bootstrap bootstrap = new Bootstrap();
        bootstrap.group(eventExecutors)
                .channel(NioSocketChannel.class)
                .option(ChannelOption.TCP_NODELAY, true)
                .handler(new ChannelInitializer<SocketChannel>() {
                    @Override
                    protected void initChannel(SocketChannel ch) throws Exception {
                        //添加服务端写客户端的Handler
                        ch.pipeline().addLast(new Dest2ClientInboundHandler(ctx));
                    }
                });
        ChannelFuture future = bootstrap.connect(msg.dstAddr(), msg.dstPort());
        future.addListener(new ChannelFutureListener() {
            @Override
            public void operationComplete(ChannelFuture future) throws Exception {
                if (future.isSuccess()) {
                    System.out.println("目标服务器连接成功");
                    //添加客户端转发请求到服务端的Handler
                    ctx.pipeline().addLast(new Client2DestInboundHandler(future));
                    DefaultSocks5CommandResponse commandResponse = new DefaultSocks5CommandResponse(Socks5CommandStatus.SUCCESS, socks5AddressType);
                    ctx.writeAndFlush(commandResponse);
                    ctx.pipeline().remove(Socks5CommandRequestInboundHandler.class);
                    ctx.pipeline().remove(Socks5CommandRequestDecoder.class);
                } else {
                    System.out.println("连接目标服务器失败,address="+msg.dstAddr()+",port=" + msg.dstPort());
                    DefaultSocks5CommandResponse commandResponse = new DefaultSocks5CommandResponse(Socks5CommandStatus.FAILURE, socks5AddressType);
                    ctx.writeAndFlush(commandResponse);
                    future.channel().close();
                }
            }
        });
    }

    private boolean inBlackList(String dstAddr) {
        for (String black : blackList) {
            if (dstAddr.toLowerCase().endsWith(black)) {
                return true;
            }
        }
        return false;
    }
}
